import asyncio
import datetime
import os
import json
import sys
import uuid
from typing import Optional, Union

from sqlalchemy import func
from sqlalchemy.orm import Session
import psutil
from fastapi import APIRouter, Depends, Request, HTTPException, Body
from fastapi.encoders import jsonable_encoder
from jose import jwt
from starlette.responses import StreamingResponse

import models
from models import Share
from models.toolkit import Cache
from schemas import response_success, response_error
from schemas.toolkit import ToolkitSchema, ToolkitUpdateSchema, ToolkitSortSchema, CacheSchema, ToolkitDebugSchema, \
    ToolkitCallSchema
from models.db import get_db
from sqlalchemy import and_, desc
from views.auth import SECRET_KEY, ALGORITHM, oauth2_scheme, get_current_user_name, token_auth, api_key_header
from views.utils.script import generate_script

toolkit_router = APIRouter()
router = toolkit_router

ws = APIRouter()

# 工具执行队列
execute_queue = {}


@router.post("/create")
def create(schema: ToolkitSchema, token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)):
    schema.username = get_current_user_name(token)

    toolkit = models.Toolkit(**jsonable_encoder(schema))
    db.add(toolkit)
    db.commit()
    return response_success("新增成功")


@router.get("/list")
def get_list(page_size: int, page: int, tool_name=None, tags=None, db: Session = Depends(get_db)):
    if not page > 0 or not page_size > 5:
        return response_error('参数错误')

    conditions = []
    # 设置查询条件
    if tool_name:
        conditions.append(models.Toolkit.name.ilike(f'%{tool_name}%'))  # 模糊匹配名称
    if tags:
        conditions.append(models.Toolkit.tags == tags)  # 精确匹配分类

    db_data = db.query(models.Toolkit).filter(and_(*conditions)).order_by(models.Toolkit.sort).all()
    print(json.dumps(jsonable_encoder(db_data), indent=4, ensure_ascii=False))
    return response_success(data={'list': db_data})


@router.post("/update")
def update(form: ToolkitUpdateSchema, db: Session = Depends(get_db)):
    if form.args:

        for args_item in form.args:
            # 处理密码类型加密
            if args_item['type'] == 'password' and len(args_item['value']) < 60:
                args_item['value'] = jwt.encode({'sub': args_item['value']}, SECRET_KEY,
                                                algorithm=ALGORITHM)

    db_data = db.query(models.Toolkit).filter_by(id=form.id).first()
    if not db_data:
        return response_error("数据不存在")

    # 不更新空数据
    for k, v in form.dict().items():
        if v is not None:
            if isinstance(v, (list, dict)):
                v = json.dumps(v)
            setattr(db_data, k, v)
    db.commit()

    return response_success("更新成功")


@router.post('/sort')
def sort(form: ToolkitSortSchema, db: Session = Depends(get_db)):
    db_data = db.query(models.Toolkit).all()
    if len(db_data) != len(form.ids):
        return response_error("请不要在筛选状态下排序")
    for item in db_data:
        item.sort = form.ids.index(item.id)

    db.commit()
    return response_success("排序成功")


@router.delete("/{id}")
def delete(id: str, db: Session = Depends(get_db)):
    db_data = db.query(models.Toolkit).filter_by(id=id).first()
    if not db_data:
        return response_error(msg="该工具不存在！")

    db.delete(db_data)
    db.commit()
    return response_success("删除成功")


@router.get("/detail")
def detail(id: str, db: Session = Depends(get_db)):
    db_data = db.query(models.Toolkit).filter_by(id=id).first()
    if db_data.args:
        db_data.args = json.loads(db_data.args)

    # data = await toolkit_collection.find_one({"_id": ObjectId(_id)})
    # data['_id'] = str(data['_id'])
    # data['id'] = str(data['_id'])
    return response_success(data={'detail': db_data})


@router.post("/call-script")
def call_script(request: Request, form: ToolkitCallSchema = Body(...), db: Session = Depends(get_db)):
    """工具调试执行"""
    token = request.headers.get("authorization")  # 获取 Authorization 头
    if token:
        token = token.replace("Bearer ", "")

    async def generate():
        db_data = db.query(models.Toolkit).filter_by(name=form.name).first()

        if not db_data:
            yield '该工具不存在 ! ...\r\n'
            return

        if not db_data.content:
            await '没有代码可供执行 ! ...\r\n'
            return

        workspace = db_data.workspace
        exec_user = db_data.exec_user
        executor = db_data.executor

        history = ''
        exit_status = -1
        if executor == 'local':
            if sys.platform == 'linux':
                from views.service.run_script_local import run_script_linux as run_script
            else:
                from views.service.run_script_local import run_script_windows as run_script

            interpreter, script = generate_script(db_data, form)
            for output in run_script(workspace=workspace, exec_user=exec_user, interpreter=interpreter, script=script):
                if isinstance(output, dict):
                    if 'history' in output:
                        history = output.pop('history')
                        exit_status = output['exit_status']
                else:
                    yield output

                await asyncio.sleep(0)

            if history:
                history = history.decode()

        elif executor == 'yunqin_terminal':
            interpreter, script = generate_script(db_data, form, False)
            from views.service.run_script_remote import run_script_yunqing

            for output in run_script_yunqing(script):
                if isinstance(output, dict):
                    exit_status = output['exit_status']
                else:
                    history += output
                    yield output
                    await asyncio.sleep(0)
        else:
            yield '该操作系统不支持 ! ...\r\n'

        if history:
            if token:
                username = get_current_user_name(token)
            else:
                username = 'SystemAPI'

            toolkit = models.ToolkitLog(tool_id=db_data.id,
                                        tool_name=db_data.name,
                                        create_time=datetime.datetime.now(),
                                        username=username,
                                        status=exit_status,
                                        content=history)
            db.add(toolkit)
            db.commit()

    return StreamingResponse(generate())


@router.post('/discontinue/{pid}')
def discontinue(pid: int):
    """终止脚本运行"""

    proces = list(filter(lambda x: x.pid == pid, psutil.process_iter(['pid'])))
    if not proces:
        return response_error(msg='没有找到终止目标')

    proces[0].kill()
    return response_success(msg='脚本终止运行')


@router.get('/history-log')
def history_log(tool_id: str, page: int, page_size: int, username: str = None, db: Session = Depends(get_db)):
    """获取工具执行历史日志"""
    if not page >= 0 or not page_size > 5:
        return response_error('参数错误')

    data = []
    # query = {
    #     'toolkit_id': toolkit_id
    # }
    # if username:
    #     query['username'] = username

    db_data = (db.query(models.ToolkitLog).filter_by(tool_id=tool_id).order_by(desc(models.ToolkitLog.create_time))
               .offset(page * page_size).limit(page_size).all())

    db_count = db.query(models.ToolkitLog).filter_by(tool_id=tool_id).count()
    return response_success(data={'list': db_data, 'total': db_count})


@router.get('/tags')
def get_tags(db: Session = Depends(get_db)):
    tags = []
    db_data = db.query(models.Toolkit).all()
    for i in db_data:
        if i.tags:
            tags.extend(i.tags.split())
    return response_success(data={'list': list(set(tags)), 'total': len(tags)})


@router.post('/cache-set')
def cache_set(form: CacheSchema, db: Session = Depends(get_db)):
    db_data = db.query(Cache).filter_by(key=form.key).first()
    if db_data:
        db_data.data = json.dumps(form.data)
    else:
        db.add(Cache(key=form.key, data=json.dumps(form.data)))

    db.commit()
    return response_success(msg='缓存成功')


@router.get('/cache-get')
def cache_get(key: str, db: Session = Depends(get_db)):
    db_data = db.query(Cache).filter_by(key=key).first()
    return response_success(data=db_data)


@router.get('/lately')
def lately(db: Session = Depends(get_db)):
    # 查询
    db_data = db.query(models.ToolkitLog).group_by(models.ToolkitLog.tool_name).order_by(
        func.max(models.ToolkitLog.create_time).desc()).limit(10).all()
    return response_success(data={'list': db_data})
